from train_jl import Train_XGBoost_TL
from predict_jl import Predict_XGBoost_TL

if __name__ == '__main__':
    # train_obj = Train_XGBoost_TL()
    # train_model = train_obj.grain_search_train_save_xgb_model()
    # train_obj.shap_graph()
    #
    # predict_obj = Predict_XGBoost_TL()
    # predict_obj.predict_xgb_model_from_localdata()
    # predict_obj.roc_graph()

    # 1.训练原始模型并生成SHAP图表
    train_obj = Train_XGBoost_TL()
    train_model = train_obj.grain_search_train_save_xgb_model()
    train_obj.shap_graph()

    # 2.自动识别并删除SHAP最小的前N个特征
    low_features = train_obj.get_low_shap_features(n_remove=3)
    reduced_model = train_obj.retrain_model_with_reduced_features(low_features)
    # reduced_model = train_obj.retrain_with_gridsearch_after_feature_removal(low_features)

    # 3.在预测文件中使用测试集评估AUC性能
    predict_obj = Predict_XGBoost_TL()
    predict_obj.compare_model_auc(
        original_model_path='../output/model/xgboost_model_jl.pkl',
        reduced_model_path='../output/model/xgboost_model_jl_reduced.pkl'
    )

    # 4.绘制RUC曲线
    # predict_obj.roc_graph()
    predict_obj.roc_graph_original_model()
    predict_obj.roc_graph_reduced_model()
    predict_obj.plot_roc_curves()


